import torch as pt
import torchvision as ptv
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

pt.manual_seed(777)

BATCH_SIZE = 16

data_dir = '../../../../large_data/DL2/pt/cifar10'
train_ds = ptv.datasets.CIFAR10(root=data_dir, train=True,
                                transform=ptv.transforms.ToTensor(),
                                download=False)
classes = train_ds.classes
test_ds = ptv.datasets.CIFAR10(root=data_dir, train=False,
                                transform=ptv.transforms.ToTensor(),
                                download=False)

train_dl = DataLoader(dataset=train_ds, batch_size=BATCH_SIZE, shuffle=True)
test_dl = DataLoader(dataset=test_ds, batch_size=BATCH_SIZE, shuffle=True)

plt.figure(figsize=[12, 6])
spr = 4
spc = 8
spn = 0


def show_pics(dl):
    global spn
    for bx, by in dl:
        for i, bxi in enumerate(bx):
            spn += 1
            if (spn > spr * spc):
                break
            plt.subplot(spr, spc, spn)
            bxi = bxi.transpose(0, 2)
            bxi = bxi.transpose(0, 1)
            plt.imshow(bxi)
            plt.axis('off')
            cls_id = by[i].item()
            plt.title(str(cls_id) + ': ' + classes[cls_id])
        if (spn > spr * spc):
            break


show_pics(train_dl)
show_pics(test_dl)
